package com.bw.test2;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.vector.VectorAssemblerBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalRegressionBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotPredictBatchOp;
import com.alibaba.alink.operator.batch.feature.OneHotTrainBatchOp;
import com.alibaba.alink.operator.batch.regression.LinearRegPredictBatchOp;
import com.alibaba.alink.operator.batch.regression.LinearRegTrainBatchOp;
import com.alibaba.alink.operator.batch.source.MemSourceBatchOp;
import com.alibaba.alink.operator.batch.utils.UDFBatchOp;
import com.alibaba.alink.operator.common.evaluation.RegressionMetrics;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.regression.LinearRegressionModel;
import org.apache.flink.table.functions.ScalarFunction;
import org.apache.flink.types.Row;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class UdfTest {

    //https://nightlies.apache.org/flink/flink-docs-release-1.20/docs/dev/table/functions/udfs/
    //https://alinklab.cn/manual/udfbatchop.html
    @Test
    public void testLinearRegTrainBatchOp() throws Exception {

        BatchOperator.setParallelism(1);
        // 1. 数据源
        List <Row> df = Arrays.asList(
                Row.of(2, "a", 1),
                Row.of(3, "b", 1),
                Row.of(4, "c", 2),
                Row.of(2, "d", 1),
                Row.of(2, "a", 1),
                Row.of(4, "b", 2),
                Row.of(1, "b", 1),
                Row.of(5, null, 3)
        );
        // 选择特征和label
        BatchOperator <?> batchData = new MemSourceBatchOp(df, "f0 int, f1 String, label int");



        UDFBatchOp udfBatchOp = new UDFBatchOp()
                .setFunc(new MyFunction())
                .setSelectedCols("f1")
                .setOutputCol("f11")
                .linkFrom(batchData);

//        udfBatchOp.print();

        String[] colnames = new String[] {"f0", "f11"};

        //  训练模型
        LinearRegression lr = new LinearRegression()
                .setFeatureCols(colnames)
                .setLabelCol("label")
                .setPredictionCol("pred")
                .enableLazyPrintModelInfo();

        LinearRegressionModel model = lr.fit(udfBatchOp);
//
//
//        // 输出结果
        model.transform(udfBatchOp).print();
    }

   public static  class MyFunction extends ScalarFunction {

        public Integer eval(String col) {
            if ("a".equals(col)){
                return 0;
            }else if ("b".equals(col)){
                return 1;
            }else if ("c".equals(col)){
                return 2;
            }else if ("d".equals(col)){
                return 3;
            }else {
                return 4;
            }
        }


    }
}